import numpy as np
from src.entity.CelestialBody import FixedStar

# 统一引力常数，可调
# 标准引力常数对于模拟来说数值过偏，不能起到很好的模拟效果
G = 2.0


class SinglyStarSystem:
    """
    单星系统：无外力，仅做匀速直线运动
    """
    def __init__(self, star: FixedStar):
        self.star = star

    def evolve(self, dt):
        # 匀速直线
        self.star.update(np.zeros(3), dt)


class DoublyStarSystem:
    """
    双星系统：两体引力
    """
    def __init__(self, s1: FixedStar, s2: FixedStar):
        self.stars = [s1, s2]

    def evolve(self, dt):
        r = self.stars[1].position - self.stars[0].position
        d = np.linalg.norm(r)
        if d == 0:
            return
        r_hat = r / d

        force = G * self.stars[0].mass * self.stars[1].mass / d**2
        # 分别计算矢量加速度
        acc0 = force * r_hat / self.stars[0].mass
        acc1 = -force * r_hat / self.stars[1].mass

        # 更新星体
        self.stars[0].update(acc0, dt)
        self.stars[1].update(acc1, dt)


class MultiStarSystem:
    """
    多星系统
    """
    def __init__(self, stars: list[FixedStar]):
        """
        :param stars 恒星列表
        """
        self.stars = stars

    def evolve(self, dt):
        """
        演化星体系统一个时间步长，通过计算星体间的万有引力加速度来更新所有星体的状态。
        使用牛顿万有引力定律计算每对星体间的引力，
        然后根据牛顿第二定律计算加速度，最后更新星体的位置和速度。

        :param dt: 时间步长，单位为秒
        """
        n = len(self.stars)
        # 初始化每个星体的加速度为零向量
        accs = [np.zeros(3) for _ in range(n)]

        # 计算所有星体对之间的引力加速度
        # 使用对称性优化计算，只计算上三角矩阵
        # 这里涉及到矩阵运算，是 线性代数 相关知识
        for i in range(n):
            for j in range(i+1, n):
                # 计算星体j相对于星体i的位置向量
                r = self.stars[j].position - self.stars[i].position
                # 计算两星体间的距离
                d = np.linalg.norm(r)
                if d == 0:
                    continue
                # 计算单位方向向量
                r_hat = r / d
                # 根据万有引力定律计算引力大小
                force = G * self.stars[i].mass * self.stars[j].mass / d**2
                # 计算加速度向量
                acc = force * r_hat
                # 根据牛顿第二定律 F=ma，将加速度分配给两个星体
                # 注意加速度方向对于不同星体是相反的
                accs[i] += acc / self.stars[i].mass
                accs[j] -= acc / self.stars[j].mass

        # 使用计算得到的加速度更新所有星体的状态
        for star, a in zip(self.stars, accs):
            star.update(a, dt)
